﻿using System;
using System.Collections.Generic;
using System.Linq;

namespace TextEmbedding.Core
{


    public class BM25Embedding : ITextEmbedding
    {
        private readonly List<string[]> _documents;
        private readonly Dictionary<string, double> _idf;
        private readonly List<string> _vocab;
        private readonly double _avgDocLength;
        private readonly double _k1;
        private readonly double _b;



        public BM25Embedding(List<string> docs, double k1 = 1.5, double b = 0.75)
        {
            _documents = docs.Select(d => ZhTokenizer.Tokenize(d)).ToList();
            _vocab = _documents.SelectMany(d => d).Distinct().ToList();
            _idf = CalculateIdf(_documents);
            _avgDocLength = _documents.Average(d => d.Length);
            _k1 = k1;
            _b = b;
        }


        private Dictionary<string, double> CalculateIdf(List<string[]> docs)
        {
            int N = docs.Count;
            var df = new Dictionary<string, int>();
            foreach (var doc in docs)
            {
                foreach (var word in doc.Distinct())
                {
                    if (!df.ContainsKey(word)) df[word] = 0;
                    df[word]++;
                }
            }

            return df.ToDictionary(kv => kv.Key, kv => Math.Log(1 + (N - kv.Value + 0.5) / (kv.Value + 0.5)));
        }

        public float[] Embedding(int dimension, string text)
        {
            var vec = new float[dimension];
            var tokens = ZhTokenizer.Tokenize(text);
            int docLength = tokens.Length;

            // 统计词频
            var tf = new Dictionary<string, int>();
            foreach (var word in tokens)
            {
                if (!tf.ContainsKey(word)) tf[word] = 0;
                tf[word]++;
            }

            for (int i = 0; i < Math.Min(dimension, _vocab.Count); i++)
            {
                string term = _vocab[i];
                if (!tf.ContainsKey(term)) continue;

                double freq = tf[term];
                double idf = _idf.ContainsKey(term) ? _idf[term] : 0;
                double score = idf * freq * (_k1 + 1) / (freq + _k1 * (1 - _b + _b * docLength / _avgDocLength));
                vec[i] = (float)score;
            }

            return vec;
        }
    }

}
